import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import seaborn as sns 

def parse_args():
    parser = argparse.ArgumentParser(description="Generate Coverage Trend Plot (Filtered).")
    parser.add_argument('--files','-f', nargs='+', required=True, help='List of files (CSV or Excel) to analyze')
    parser.add_argument('--exclude', nargs='+', default=['BART'], help='Methods to exclude (default: BART)')
    parser.add_argument('--output','-o', default='coverage_trend_filtered.png', help='Output filename for the plot')
    return parser.parse_args()

def read_data(file_path):
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        else:
            return pd.read_csv(file_path)
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

def get_aggregated_trend(files, exclude_methods):
    """
    读取所有文件，过滤非整10的Coverage，合并后计算平均值。
    """
    all_data = []
    
    for f in files:
        df = read_data(f)
        if df is not None and not df.empty:
            # 1. 过滤不需要的方法
            df = df[~df['method'].isin(exclude_methods)].copy()
            
            # 2. === 重命名 ===
            df['method'] = df['method'].replace('ours(MTM)', 'ours(deepsignal3)')
            
            # 3. 确保 coverage 是数值型
            df['coverage'] = pd.to_numeric(df['coverage'], errors='coerce')
            
            # 4. === 核心修改：只保留 10 的整数倍 Coverage ===
            # 使用 np.isclose 避免浮点数精度问题 (例如 30.00000001)
            # 逻辑：余数接近0
            df = df[np.isclose(df['coverage'] % 10, 0)]
            
            # 只保留需要的列
            cols = ['method', 'coverage', 'pearson', 'rsquare', 'spearman', 'RMSE']
            df = df[[c for c in cols if c in df.columns]]
            all_data.append(df)
            
    if not all_data:
        return None
        
    combined_df = pd.concat(all_data, ignore_index=True)
    
    # 5. 聚合：计算每个方法在每个Coverage下的平均表现
    trend_df = combined_df.groupby(['method', 'coverage']).mean().reset_index()
    
    # 排序：确保绘图时线条是按coverage顺序连起来的
    trend_df = trend_df.sort_values(by=['method', 'coverage'])
    
    return trend_df

def plot_trend(trend_df, output_file):
    # 设置风格
    try:
        sns.set_style("whitegrid")
        sns.set_context("paper", font_scale=1.2)
    except:
        plt.style.use('ggplot')

    # 定义四个子图的配置
    metrics = [
        {'col': 'pearson', 'title': 'Pearson Correlation (Higher is Better)'},
        {'col': 'rsquare', 'title': 'R-Square (Higher is Better)'},
        {'col': 'spearman', 'title': 'Spearman Correlation (Higher is Better)'},
        {'col': 'RMSE', 'title': 'RMSE (Lower is Better)'} 
    ]
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    axes = axes.flatten()
    
    methods = trend_df['method'].unique()
    
    # === 样式定义 ===
    # Dorado: 蓝色实线 (基准)
    # ours(deepsignal3): 红色加粗实线 (重点)
    # Rockfish/DeepMod2: 虚线 (背景对比)
    style_map = {
        'Dorado': {'color': '#1f77b4', 'ls': '-', 'lw': 2, 'marker': 'o', 'zorder': 5},
        'ours(deepsignal3)': {'color': '#d62728', 'ls': '-', 'lw': 3, 'marker': 'D', 'zorder': 10}, 
        'Rockfish': {'color': '#2ca02c', 'ls': '--', 'lw': 2, 'marker': '^', 'zorder': 3},
        'DeepMod2': {'color': '#ff7f0e', 'ls': '-.', 'lw': 2, 'marker': 's', 'zorder': 3}
    }
    
    for i, config in enumerate(metrics):
        ax = axes[i]
        metric_col = config['col']
        
        for method in methods:
            subset = trend_df[trend_df['method'] == method]
            
            # 获取样式
            s = style_map.get(method, {'color': 'gray', 'ls': ':', 'lw': 1, 'marker': '.', 'zorder': 1})
            
            ax.plot(subset['coverage'], subset[metric_col], 
                    label=method,
                    color=s['color'], 
                    linestyle=s['ls'], 
                    linewidth=s['lw'],
                    marker=s['marker'],
                    markersize=6,
                    zorder=s['zorder'])
        
        ax.set_title(config['title'], fontweight='bold')
        ax.set_xlabel('Coverage (X)')
        ax.set_ylabel(metric_col)
        
        # 强制 X 轴只显示存在的 Coverage 点 (如 10, 20, 30...)
        unique_coverages = sorted(trend_df['coverage'].unique())
        if len(unique_coverages) > 0:
            ax.set_xticks(unique_coverages)
        
        sns.despine(ax=ax)

    # 统一图例
    handles, labels = axes[0].get_legend_handles_labels()
    # 调整图例位置到图片最下方
    fig.legend(handles, labels, loc='lower center', bbox_to_anchor=(0.5, -0.02), ncol=4, fontsize=12, frameon=False)
    
    plt.tight_layout()
    plt.subplots_adjust(bottom=0.08) # 留出空间给图例
    
    print(f"Saving filtered trend plot to {output_file}")
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.show()

def main():
    args = parse_args()
    
    print("Reading and aggregating data...")
    trend_df = get_aggregated_trend(args.files, args.exclude)
    
    if trend_df is not None:
        print(f"Data loaded. Processing coverages: {sorted(trend_df['coverage'].unique())}")
        plot_trend(trend_df, args.output)
    else:
        print("No data found (check if files contain coverages 10, 20, 30...).")

if __name__ == "__main__":
    main()